#dgl.graph可视化-按头尾结点列表和边权列表绘图
#如何仅绘制前多少个结点？如何由字典绘图？
import dgl
import torch as th
import networkx as nx
import matplotlib.pyplot as plt

u = th.tensor([0, 1, 2, 3])
v = th.tensor([1, 2, 3, 0])
g = dgl.graph((u, v))  # u-头结点列表，v-尾节点列表，uv按列对应
g.ndata["desc"] = g.nodes()  # 节点下标，显示N:0，使用g.ndata[‘x’]可以为节点添加特征或访问节点特征，数量要和g.nodes()对应
g.edata["weights"] = th.tensor([4, 5, 6, 7])  # 边权，使用g.ndata[‘x’]可以为节点添加特征或访问节点特征，数量要和g.edges()对应


# # 如果有孤立节点，需要在创建边时显示设置节点数量g = dgl.graph((u, v), num_nodes=8)
# print(g.edges())#边(tensor([0, 1, 2, 3]), tensor([2, 0, 1, 5]))
# print(g.nodes())#节点列表（从0开始，不和u、v值一样）tensor([0, 1, 2, 3, 4, 5])
# #头节点列表和尾节点列表，还有边的IDs
# print(g.edges(form='all'))#(tensor([0, 1, 2, 3]), tensor([2, 0, 1, 5]), tensor([0, 1, 2, 3]))
# #将单向图转化为双向图
# bg = dgl.to_bidirected(g)
# print(bg.edges())#(tensor([0, 0, 1, 1, 2, 2, 3, 5]), tensor([1, 2, 0, 2, 0, 1, 5, 3]))

# desc, weights
#
# 图形形状会随机调整
def ShowGraph(graph, nodeLabel, EdgeLabel):
    plt.figure(figsize=(8, 8))
    G = graph.to_networkx(node_attrs=nodeLabel.split(), edge_attrs=EdgeLabel.split())  # 转换 dgl graph to networks
    pos = nx.spring_layout(G)
    nx.draw(G, pos, edge_color="grey", node_size=500, with_labels=True)  # 画图，设置节点大小
    node_data = nx.get_node_attributes(G, nodeLabel)  # 获取节点的desc属性
    node_labels = {index: "N:" + str(data) for index, data in
                   enumerate(node_data)}  # 重新组合数据， 节点标签是dict, {nodeid:value,nodeid2,value2} 这样的形式
    pos_higher = {}

    for k, v in pos.items():  # 调整下顶点属性显示的位置，不要跟顶点的序号重复了
        if (v[1] > 0):
            pos_higher[k] = (v[0] - 0.04, v[1] + 0.04)
        else:
            pos_higher[k] = (v[0] - 0.04, v[1] - 0.04)
    nx.draw_networkx_labels(G, pos_higher, labels=node_labels, font_color="brown", font_size=12)  # 将desc属性，显示在节点上
    edge_labels = nx.get_edge_attributes(G, EdgeLabel)  # 获取边的weights属性，

    edge_labels = {(key[0], key[1]): "w:" + str(edge_labels[key].item()) for key in
                   edge_labels}  # 重新组合数据， 边的标签是dict, {(nodeid1,nodeid2):value,...} 这样的形式
    nx.draw_networkx_edges(G, pos, alpha=0.5)
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=12)  # 将Weights属性，显示在边上

    print(G.edges.data())
    # plt.show()
    save_fig_name="graph_with_"+str(len(G.nodes()))+"nodes"
    plt.savefig(save_fig_name)#保存图片

ShowGraph(g, "desc", 'weights')
plt.close()
